# -*- coding: utf-8 -*-
"""
彗星亮度预测工具 v8.1
- 支持彗星和小行星
- 根据选择自动调整输入字段
- 修复输入错误
- 新增数据自动识别功能
- 新增导出数据到CSV功能
- 新增观测数据导入功能
- 在亮度图表中显示观测数据
- 改进观测数据导入功能
- UI美化
"""
import threading
import re
import csv
import os
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from tkinter import ttk, messagebox, scrolledtext, filedialog
import tkinter as tk
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import numpy as np
import matplotlib
matplotlib.use('TkAgg')

# ---------- 中文显示 ----------
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# ---------- 样式配置 ----------


class StyleConfig:
    PRIMARY = "#2c3e50"
    SECONDARY = "#34495e"
    ACCENT = "#3498db"
    SUCCESS = "#27ae60"
    WARNING = "#f39c12"
    DANGER = "#e74c3c"
    LIGHT = "#ecf0f1"
    DARK = "#2c3e50"
    FONT_FAMILY = "微软雅黑"
    FONT_SIZE = 10

# ---------- 美化组件 ----------


class ModernButton(ttk.Button):
    def __init__(self, master=None, **kwargs):
        style = ttk.Style()
        style.configure("Modern.TButton",
                        font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE),
                        padding=(10, 5))
        kwargs["style"] = "Modern.TButton"
        super().__init__(master, **kwargs)


class ModernEntry(ttk.Entry):
    def __init__(self, master=None, **kwargs):
        style = ttk.Style()
        style.configure("Modern.TEntry",
                        font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE-1))
        kwargs["style"] = "Modern.TEntry"
        super().__init__(master, **kwargs)


class ModernLabel(ttk.Label):
    def __init__(self, master=None, **kwargs):
        style = ttk.Style()
        style.configure("Modern.TLabel",
                        font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE),
                        background=StyleConfig.LIGHT)
        kwargs["style"] = "Modern.TLabel"
        super().__init__(master, **kwargs)


class ModernFrame(ttk.Frame):
    def __init__(self, master=None, **kwargs):
        style = ttk.Style()
        style.configure("Modern.TFrame", background=StyleConfig.LIGHT)
        kwargs["style"] = "Modern.TFrame"
        super().__init__(master, **kwargs)

# ---------- 主程序 ----------


class CometBrightnessApp:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("彗星与小行星亮度预测工具 v8.1")
        self.root.geometry("1400x900")
        self.root.minsize(1200, 700)

        # 设置主题
        self.setup_theme()

        # 设置图标（如果有的话）
        try:
            self.root.iconbitmap("comet.ico")  # 如果有图标文件的话
        except:
            pass

        # 数据缓存
        self.current_dates = None
        self.cobs_calc_m_pred = None
        self.cobs_obs_m_pred = None
        self.astro_m_pred = None
        self.mpc_m_pred = None
        self.r_pred = None
        self.delta_pred = None

        # 观测数据缓存
        self.observation_data = {
            'comet': [],  # 彗星观测数据
            'asteroid': []  # 小行星观测数据
        }

        # 子窗口引用
        self.brightness_window = None
        self.distance_window = None

        # 进度条变量
        self.progress_var = tk.IntVar(value=0)
        self.progress_text = tk.StringVar(value="就绪")

        self._build_ui()
        self.root.after(100, self.fill_sample_data)

    def setup_theme(self):
        """设置现代化主题"""
        style = ttk.Style()

        # 尝试设置现代化主题
        try:
            style.theme_use('clam')
        except:
            try:
                style.theme_use('alt')
            except:
                pass

        # 配置各种样式
        style.configure("Modern.TFrame", background=StyleConfig.LIGHT)
        style.configure("Modern.TLabel", background=StyleConfig.LIGHT,
                        font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE))
        style.configure("Modern.TButton", font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE),
                        padding=(10, 5))
        style.configure("Modern.TEntry", font=(
            StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE-1))
        style.configure("Custom.TNotebook", background=StyleConfig.LIGHT)
        style.configure("Custom.TNotebook.Tab", font=(StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE),
                        padding=(10, 5))
        style.configure("Header.TLabel", font=(StyleConfig.FONT_FAMILY, 12, "bold"),
                        background=StyleConfig.PRIMARY, foreground="white")
        style.configure("Success.TButton", background=StyleConfig.SUCCESS)
        style.configure("Primary.TButton", background=StyleConfig.PRIMARY)
        style.configure("Accent.TButton", background=StyleConfig.ACCENT)

    # -------------------- UI 构建 --------------------
    def _build_ui(self):
        # 创建主容器
        main_container = ModernFrame(self.root)
        main_container.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # 标题
        title_frame = ModernFrame(main_container)
        title_frame.pack(fill=tk.X, pady=(0, 10))

        title_label = ModernLabel(title_frame, text="彗星与小行星亮度预测工具 v8.1",
                                  style="Header.TLabel")
        title_label.pack(fill=tk.X, padx=10, pady=10)

        main_paned = ttk.PanedWindow(main_container, orient=tk.HORIZONTAL)
        main_paned.pack(fill=tk.BOTH, expand=True)

        left_frame = ModernFrame(main_paned, width=450)
        main_paned.add(left_frame, weight=0)

        right_frame = ModernFrame(main_paned)
        main_paned.add(right_frame, weight=1)

        # 左侧：参数 + 控制
        self._build_left(left_frame)

        # 右侧：结果 + 进度
        self._build_right(right_frame)

    def _build_left(self, master):
        # 使用自定义样式的Notebook
        style = ttk.Style()
        style.configure("Custom.TNotebook", background=StyleConfig.LIGHT)
        style.configure("Custom.TNotebook.Tab", font=(
            StyleConfig.FONT_FAMILY, StyleConfig.FONT_SIZE))

        notebook = ttk.Notebook(master, style="Custom.TNotebook")
        notebook.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 各页面
        pages = [
            ("基本参数", self._build_basic_page),
            ("模型参数", self._build_model_page),
            ("预测控制", self._build_control_page),
            ("轨道数据导入", self._build_import_page),
            ("观测数据导入", self._build_observation_page)
        ]

        for page_name, page_builder in pages:
            frame = ModernFrame(notebook)
            notebook.add(frame, text=page_name)
            page_builder(frame)

    def _build_basic_page(self, master):
        # 彗星/小行星选择
        comet_frame = ttk.LabelFrame(
            master, text="天体选择", padding=15, style="Modern.TLabelframe")
        comet_frame.pack(fill=tk.X, pady=5)

        ModernLabel(comet_frame, text="选择天体类型:").grid(
            row=0, column=0, sticky=tk.W, padx=5, pady=5)
        self.comet_type_var = tk.StringVar()
        cb = ttk.Combobox(
            comet_frame, textvariable=self.comet_type_var, width=25, state="readonly")
        cb["values"] = ["彗星", "小行星"]
        cb.grid(row=0, column=1, padx=5, pady=5)
        cb.bind("<<ComboboxSelected>>", self.on_comet_type_selected)

        # 天体信息
        info_frame = ttk.LabelFrame(master, text="天体信息", padding=15)
        info_frame.pack(fill=tk.X, pady=5)

        ModernLabel(info_frame, text="天体名称:").grid(
            row=0, column=0, sticky=tk.W, padx=5, pady=5)
        self.name_entry = ModernEntry(info_frame, width=25)
        self.name_entry.grid(row=0, column=1, padx=5, pady=5)

        # 轨道元素
        orbit_frame = ttk.LabelFrame(master, text="轨道元素", padding=15)
        orbit_frame.pack(fill=tk.X, pady=5)

        labels = [
            ("q (近日点距离 AU)", "perihelion_distance"),
            ("a (半长轴 AU)", "semi_major_axis"),
            ("e (偏心率)", "eccentricity"),
            ("ω (近日点幅角 °)", "arg_perihelion"),
            ("Ω (升交点赤经 °)", "long_asc_node"),
            ("i (轨道倾角 °)", "inclination"),
            ("P (轨道周期 年)", "orbital_period"),
            ("T (近日点时间)", "perihelion_time"),
            ("历元时间", "epoch_time")
        ]

        self.orbit_entries = {}
        for idx, (txt, key) in enumerate(labels):
            ModernLabel(orbit_frame, text=txt).grid(
                row=idx, column=0, sticky=tk.W, padx=5, pady=2)
            ent = ModernEntry(orbit_frame, width=18)
            ent.grid(row=idx, column=1, padx=5, pady=2)
            self.orbit_entries[key] = ent

    def _build_model_page(self, master):
        # 模型开关
        switch_frame = ttk.LabelFrame(master, text="模型开关", padding=15)
        switch_frame.pack(fill=tk.X, pady=5)

        self.use_cobs_calc_var = tk.BooleanVar(value=True)
        self.use_cobs_obs_var = tk.BooleanVar(value=True)
        self.use_astro_var = tk.BooleanVar(value=True)
        self.use_mpc_var = tk.BooleanVar(value=True)

        ttk.Checkbutton(switch_frame, text="启用 MPC 模型",
                        variable=self.use_mpc_var).pack(anchor=tk.W, pady=2)
        ttk.Checkbutton(switch_frame, text="启用 COBS-Calc 模型",
                        variable=self.use_cobs_calc_var).pack(anchor=tk.W, pady=2)
        ttk.Checkbutton(switch_frame, text="启用 COBS-Observe 模型",
                        variable=self.use_cobs_obs_var).pack(anchor=tk.W, pady=2)
        ttk.Checkbutton(switch_frame, text="启用 ASTRO 模型",
                        variable=self.use_astro_var).pack(anchor=tk.W, pady=2)

        # 星等参数
        mag_frame = ttk.LabelFrame(master, text="星等参数 (H/n)", padding=15)
        mag_frame.pack(fill=tk.X, pady=5)

        self.cobs_calc_H_entry, self.cobs_calc_n_entry = self._add_hn_row(
            mag_frame, 0, "COBS-Calc")
        self.cobs_obs_H_entry, self.cobs_obs_n_entry = self._add_hn_row(
            mag_frame, 1, "COBS-Observe")
        self.astro_H_entry, self.astro_n_entry = self._add_hn_row(
            mag_frame, 2, "ASTRO")
        self.mpc_H_entry, self.mpc_n_entry = self._add_hn_row(
            mag_frame, 3, "MPC")

    def _add_hn_row(self, master, row, label):
        ModernLabel(master, text=f"{label} H:").grid(
            row=row, column=0, padx=5, pady=2, sticky=tk.W)
        h = ModernEntry(master, width=8)
        h.grid(row=row, column=1, padx=5, pady=2)
        ModernLabel(master, text="n:").grid(row=row, column=2, padx=5, pady=2)
        n = ModernEntry(master, width=8)
        n.grid(row=row, column=3, padx=5, pady=2)
        return h, n

    def _build_control_page(self, master):
        # 预测模式
        mode_frame = ttk.LabelFrame(master, text="预测模式", padding=15)
        mode_frame.pack(fill=tk.X, pady=5)

        self.pred_mode = tk.StringVar(value="days")
        ttk.Radiobutton(mode_frame, text="按天数预测", variable=self.pred_mode,
                        value="days").pack(anchor=tk.W, pady=2)
        ttk.Radiobutton(mode_frame, text="按日期范围预测", variable=self.pred_mode,
                        value="dates").pack(anchor=tk.W, pady=2)
        self.pred_mode.trace_add("write", self.toggle_pred_mode)

        self.days_frame = ModernFrame(mode_frame)
        self.days_frame.pack(fill=tk.X, pady=5)
        ModernLabel(self.days_frame, text="预测天数:").pack(side=tk.LEFT, padx=5)
        self.pred_days_entry = ModernEntry(self.days_frame, width=10)
        self.pred_days_entry.insert(0, "365")
        self.pred_days_entry.pack(side=tk.LEFT, padx=5)

        self.dates_frame = ModernFrame(mode_frame)
        ModernLabel(self.dates_frame, text="开始日期:").pack(side=tk.LEFT, padx=5)
        self.start_date_entry = ModernEntry(self.dates_frame, width=12)
        self.start_date_entry.pack(side=tk.LEFT, padx=5)
        ModernLabel(self.dates_frame, text="结束日期:").pack(side=tk.LEFT, padx=5)
        self.end_date_entry = ModernEntry(self.dates_frame, width=12)
        self.end_date_entry.pack(side=tk.LEFT, padx=5)

        # 按钮区域
        btn_frame = ModernFrame(master)
        btn_frame.pack(fill=tk.X, pady=15)

        ModernButton(btn_frame, text="开始预测", command=self.run_prediction,
                     style="Success.TButton").pack(side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="重置参数", command=self.clear_entries).pack(
            side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="帮助", command=self.show_help).pack(
            side=tk.RIGHT, padx=5)

    def _build_import_page(self, master):
        """构建轨道数据导入页面"""
        import_frame = ttk.LabelFrame(master, text="轨道数据自动识别导入", padding=15)
        import_frame.pack(fill=tk.BOTH, expand=True, pady=5)

        # 说明文本
        help_text = """支持从COBS或MPC网站复制轨道数据自动识别：
彗星数据示例 (COBS格式):
q (perihelion distance)    0.504093
a (semi-major axes)        84.678817
e (eccentricity)           0.994047
ω (argument of perihelion) 308.3520
Ω (Longitude of ascending node) 335.3310
i (inclination)            4.4728
P (Orbital period in years) 779.2551
T (time of perihelion passage) 2025-09-12 20:12:11
Epoch 2025 Nov 21
Absolute magnitude         8.9
Slope parameter            5.7
"""
        help_label = ModernLabel(import_frame, text=help_text, justify=tk.LEFT)
        help_label.pack(fill=tk.X, pady=5)

        # 数据输入文本框
        ModernLabel(import_frame, text="粘贴轨道数据:").pack(
            anchor=tk.W, pady=(10, 5))
        self.data_text = scrolledtext.ScrolledText(
            import_frame, height=10, width=50, font=(StyleConfig.FONT_FAMILY, 9))
        self.data_text.pack(fill=tk.BOTH, expand=True, pady=5)

        # 按钮框架
        btn_frame = ModernFrame(import_frame)
        btn_frame.pack(fill=tk.X, pady=10)

        ModernButton(btn_frame, text="识别并导入轨道数据", command=self.auto_import_data,
                     style="Primary.TButton").pack(side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="清空数据", command=self.clear_import_data).pack(
            side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="导入示例数据", command=self.import_sample_data).pack(
            side=tk.LEFT, padx=5)

    def _build_observation_page(self, master):
        """构建观测数据导入页面"""
        obs_frame = ttk.LabelFrame(master, text="观测数据导入", padding=15)
        obs_frame.pack(fill=tk.BOTH, expand=True, pady=5)

        # 说明文本
        help_text = """支持导入观测数据文件：

1. 彗星观测数据 (CSV格式):
   - 支持COBS导出的CSV格式
   - 包含"Obs date"和"Mag"列
   - 示例: "2025-11-18 18:14", "18.6"

2. 小行星观测数据 (TXT格式):
   - 支持MPC的80列格式
   - 包含日期和星等信息
   - 示例: "2023 02 15.70243", "8.2 V"

导入后观测数据将在亮度图表中显示为散点。
"""
        help_label = ModernLabel(obs_frame, text=help_text, justify=tk.LEFT)
        help_label.pack(fill=tk.X, pady=5)

        # 按钮框架
        btn_frame = ModernFrame(obs_frame)
        btn_frame.pack(fill=tk.X, pady=10)

        ModernButton(btn_frame, text="导入彗星观测数据(CSV)", command=self.import_comet_observations,
                     style="Accent.TButton").pack(side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="导入小行星观测数据(TXT)", command=self.import_asteroid_observations,
                     style="Accent.TButton").pack(side=tk.LEFT, padx=5)
        ModernButton(btn_frame, text="清空观测数据", command=self.clear_observation_data).pack(
            side=tk.LEFT, padx=5)

        # 观测数据显示区域
        ModernLabel(obs_frame, text="当前观测数据:").pack(anchor=tk.W, pady=(10, 5))
        self.obs_text = scrolledtext.ScrolledText(
            obs_frame, height=8, width=50, font=(StyleConfig.FONT_FAMILY, 9))
        self.obs_text.pack(fill=tk.BOTH, expand=True, pady=5)
        self.obs_text.config(state=tk.DISABLED)

    def _build_right(self, master):
        top = ModernFrame(master)
        top.pack(fill=tk.BOTH, expand=True)

        self._build_result(top)

        bottom = ModernFrame(master, height=40)
        bottom.pack(fill=tk.X, side=tk.BOTTOM, padx=5, pady=2)

        # 进度条区域
        progress_frame = ModernFrame(bottom)
        progress_frame.pack(fill=tk.X, pady=5)

        ModernLabel(progress_frame, text="进度:").pack(side=tk.LEFT, padx=(0, 5))
        self.progress_bar = ttk.Progressbar(
            progress_frame, variable=self.progress_var, maximum=100)
        self.progress_bar.pack(fill=tk.X, side=tk.LEFT,
                               expand=True, padx=(0, 8))
        ModernLabel(progress_frame, textvariable=self.progress_text,
                    width=20).pack(side=tk.RIGHT)

        # 图表控制区域
        chart = ttk.LabelFrame(top, text="图表窗口", padding=15)
        chart.pack(fill=tk.X, pady=5)

        # 导出按钮
        export_frame = ModernFrame(chart)
        export_frame.pack(fill=tk.X, pady=5)
        ModernButton(export_frame, text="导出预测数据到CSV", command=self.export_to_csv,
                     style="Success.TButton").pack(side=tk.LEFT, padx=5)

        # 图表控制按钮
        chart_btn_frame = ModernFrame(chart)
        chart_btn_frame.pack(fill=tk.X, pady=5)

        for txt, cmd in [("打开亮度图表", self.open_brightness_window),
                         ("打开距离图表", self.open_distance_window),
                         ("关闭所有图表", self.close_all_windows)]:
            ModernButton(chart_btn_frame, text=txt,
                         command=cmd).pack(side=tk.LEFT, padx=3)

    def _build_result(self, master):
        result = ttk.LabelFrame(master, text="预测结果", padding=15)
        result.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 创建滚动区域
        canvas = tk.Canvas(result, bg=StyleConfig.LIGHT)
        scroll = ttk.Scrollbar(result, orient="vertical", command=canvas.yview)
        scrollable = ModernFrame(canvas)

        scrollable.bind("<Configure>", lambda e: canvas.configure(
            scrollregion=canvas.bbox("all")))
        canvas.create_window((0, 0), window=scrollable, anchor="nw")
        canvas.configure(yscrollcommand=scroll.set)

        canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scroll.pack(side=tk.RIGHT, fill=tk.Y)

        # 结果项
        items = [
            "轨道类型", "近日点日期", "预测总天数", "近地点距离",
            "COBS-Calc最亮星等", "COBS-Calc最亮日期",
            "COBS-Observe最亮星等", "COBS-Observe最亮日期",
            "ASTRO最亮星等", "ASTRO最亮日期",
            "MPC最亮星等", "MPC最亮日期",
            "今日日彗距", "今日地彗距",
            "今日COBS-Calc星等", "今日COBS-Observe星等", "今日ASTRO星等", "今日MPC星等",
            "30天后日彗距", "30天后地彗距", "30天后COBS-Calc星等", "30天后COBS-Observe星等", "30天后ASTRO星等", "30天后MPC星等",
            "60天后日彗距", "60天后地彗距", "60天后COBS-Calc星等", "60天后COBS-Observe星等", "60天后ASTRO星等", "60天后MPC星等",
            "彗发长度 (km)", "彗发长度 (°)", "彗尾长度 (km)", "彗尾长度 (°)"
        ]

        self.info_labels = {}
        for k in items:
            f = ModernFrame(scrollable)
            f.pack(fill=tk.X, pady=1)
            ModernLabel(f, text=k+":", width=26,
                        anchor="e").pack(side=tk.LEFT, padx=5)
            lab = ModernLabel(f, text="待计算", width=41, anchor="w")
            lab.pack(side=tk.LEFT, padx=5)
            self.info_labels[k] = lab

    def show_help(self):
        """显示帮助信息"""
        help_text = """
彗星与小行星亮度预测工具 v8.1 使用说明

基本功能:
1. 选择天体类型（彗星或小行星）
2. 输入轨道参数和星等参数
3. 选择预测模式（按天数或日期范围）
4. 点击开始预测进行计算

数据导入:
- 轨道数据导入：支持从COBS/MPC网站复制数据自动识别
- 观测数据导入：支持导入历史观测数据并在图表中显示

图表功能:
- 亮度预测图表：显示各模型的亮度预测曲线
- 距离变化图表：显示日彗距和地彗距变化
- 支持导出数据到CSV文件

模型说明:
- MPC模型：适用于小行星
- COBS模型：适用于彗星
- ASTRO模型：通用模型

提示: 首次使用可点击"导入示例数据"快速体验
        """
        messagebox.showinfo("帮助", help_text)

    # -------------------- 观测数据导入功能 --------------------
    def import_comet_observations(self):
        """导入彗星观测数据"""
        filename = filedialog.askopenfilename(
            title="选择彗星观测数据文件",
            filetypes=[("CSV文件", "*.csv"), ("所有文件", "*.*")]
        )

        if not filename:
            return

        try:
            with open(filename, 'r', encoding='utf-8') as file:
                reader = csv.DictReader(file)
                observations = []

                for row in reader:
                    # 解析日期和星等
                    if 'Obs date' in row and 'Mag' in row:
                        try:
                            obs_date = datetime.strptime(
                                row['Obs date'].strip(), "%Y-%m-%d %H:%M")
                            magnitude = float(row['Mag'].strip())

                            observations.append({
                                'date': obs_date,
                                'magnitude': magnitude,
                                'type': 'comet',
                                'source': 'COBS'
                            })
                        except (ValueError, KeyError):
                            continue

                self.observation_data['comet'] = observations
                self._update_observation_display()
                messagebox.showinfo("成功", f"成功导入 {len(observations)} 条彗星观测数据")

        except Exception as e:
            messagebox.showerror("导入错误", f"导入彗星观测数据失败: {str(e)}")

    def import_asteroid_observations(self):
        """导入小行星观测数据"""
        filename = filedialog.askopenfilename(
            title="选择小行星观测数据文件",
            filetypes=[("文本文件", "*.txt"), ("所有文件", "*.*")]
        )

        if not filename:
            return

        try:
            with open(filename, 'r', encoding='utf-8') as file:
                observations = []

                for line in file:
                    line = line.strip()
                    if len(line) >= 71:  # 确保是有效的80列格式
                        try:
                            # 解析日期 (列16-32) 格式: "YYYY MM DD.dddddd"
                            date_str = line[15:32].strip()

                            # 分割日期部分
                            parts = date_str.split()
                            if len(parts) >= 3:
                                year = int(parts[0])
                                month = int(parts[1])
                                day_frac = float(parts[2])
                                day = int(day_frac)

                                # 计算小数部分对应的时间
                                fractional_day = day_frac - day
                                total_seconds = fractional_day * 86400  # 一天的秒数
                                hours = int(total_seconds // 3600)
                                minutes = int((total_seconds % 3600) // 60)

                                obs_date = datetime(
                                    year, month, day, hours, minutes)

                                # 解析星等 (列66-71) 格式: "XX.XXc" 或 "XX.XX V" 等
                                mag_str = line[65:71].strip()

                                # 提取数字部分
                                magnitude_match = re.search(
                                    r'(\d+\.\d+)', mag_str)
                                if magnitude_match:
                                    magnitude = float(magnitude_match.group(1))

                                    observations.append({
                                        'date': obs_date,
                                        'magnitude': magnitude,
                                        'type': 'asteroid',
                                        'source': 'MPC'
                                    })

                        except (ValueError, IndexError, Exception) as e:
                            continue

            if observations:
                self.observation_data['asteroid'] = observations
                self._update_observation_display()
                messagebox.showinfo("成功", f"成功导入 {len(observations)} 条小行星观测数据")
            else:
                messagebox.showwarning("警告", "未找到有效的小行星观测数据")

        except Exception as e:
            messagebox.showerror("导入错误", f"导入小行星观测数据失败: {str(e)}")

    def clear_observation_data(self):
        """清空观测数据"""
        self.observation_data = {'comet': [], 'asteroid': []}
        self._update_observation_display()
        messagebox.showinfo("成功", "已清空所有观测数据")

    def _update_observation_display(self):
        """更新观测数据显示"""
        self.obs_text.config(state=tk.NORMAL)
        self.obs_text.delete("1.0", tk.END)

        total_comet = len(self.observation_data['comet'])
        total_asteroid = len(self.observation_data['asteroid'])

        self.obs_text.insert(tk.END, f"彗星观测数据: {total_comet} 条\n")
        self.obs_text.insert(tk.END, f"小行星观测数据: {total_asteroid} 条\n\n")

        # 显示数据详情
        if total_asteroid > 0:
            self.obs_text.insert(tk.END, "小行星观测数据详情:\n")
            # 只显示前10条
            for i, obs in enumerate(self.observation_data['asteroid'][:10]):
                self.obs_text.insert(
                    tk.END, f"  {i+1}. {obs['date'].strftime('%Y-%m-%d %H:%M')}: {obs['magnitude']:.2f}等\n")
            if total_asteroid > 10:
                self.obs_text.insert(
                    tk.END, f"  ... 还有 {total_asteroid - 10} 条数据\n")

        if total_comet > 0:
            self.obs_text.insert(tk.END, "\n彗星观测数据详情:\n")
            for i, obs in enumerate(self.observation_data['comet'][:5]):
                self.obs_text.insert(
                    tk.END, f"  {i+1}. {obs['date'].strftime('%Y-%m-%d %H:%M')}: {obs['magnitude']:.2f}等\n")

        self.obs_text.config(state=tk.DISABLED)

    # -------------------- 数据导出功能 --------------------
    def export_to_csv(self):
        """导出数据到CSV文件"""
        if self.current_dates is None:
            messagebox.showwarning("警告", "请先进行预测计算再导出数据")
            return

        filename = filedialog.asksaveasfilename(
            title="保存CSV文件",
            defaultextension=".csv",
            filetypes=[("CSV文件", "*.csv"), ("所有文件", "*.*")],
            initialfile=f"{self.name_entry.get()}_亮度预测.csv"
        )

        if not filename:
            return

        try:
            with open(filename, 'w', newline='', encoding='utf-8-sig') as csvfile:
                writer = csv.writer(csvfile)

                # 写入表头
                headers = ["天体名称", "时间", "MPC星等", "COBS-Calc星等",
                           "COBS-Observe星等", "ASTRO星等", "彗日距", "彗地距"]
                writer.writerow(headers)

                # 写入数据
                for i, date in enumerate(self.current_dates):
                    date_str = date.strftime("%Y-%m-%d %H:%M")

                    mpc_mag = f"{self.mpc_m_pred[i]:.2f}" if self.mpc_m_pred is not None else ""
                    cobs_calc_mag = f"{self.cobs_calc_m_pred[i]:.2f}" if self.cobs_calc_m_pred is not None else ""
                    cobs_obs_mag = f"{self.cobs_obs_m_pred[i]:.2f}" if self.cobs_obs_m_pred is not None else ""
                    astro_mag = f"{self.astro_m_pred[i]:.2f}" if self.astro_m_pred is not None else ""

                    r = f"{self.r_pred[i]:.6f}" if self.r_pred is not None else ""
                    delta = f"{self.delta_pred[i]:.6f}" if self.delta_pred is not None else ""

                    row = [
                        self.name_entry.get(),
                        date_str,
                        mpc_mag,
                        cobs_calc_mag,
                        cobs_obs_mag,
                        astro_mag,
                        r,
                        delta
                    ]
                    writer.writerow(row)

            messagebox.showinfo("导出成功", f"数据已成功导出到:\n{filename}")

        except Exception as e:
            messagebox.showerror("导出错误", f"导出数据时发生错误:\n{str(e)}")

    # -------------------- 轨道数据导入功能 --------------------
    def auto_import_data(self):
        """自动识别并导入轨道数据"""
        data = self.data_text.get("1.0", tk.END).strip()
        if not data:
            messagebox.showwarning("警告", "请输入要识别的数据")
            return

        try:
            if self._detect_and_import_asteroid_data(data):
                messagebox.showinfo("成功", "小行星数据导入成功！")
            elif self._detect_and_import_comet_data(data):
                messagebox.showinfo("成功", "彗星数据导入成功！")
            else:
                messagebox.showerror("错误", "无法识别数据格式，请检查数据格式")
        except Exception as e:
            messagebox.showerror("导入错误", f"数据导入失败: {str(e)}")

    def _detect_and_import_comet_data(self, data):
        """检测并导入彗星轨道数据"""
        lines = data.split('\n')

        comet_keywords = ['perihelion distance', 'semi-major axes', 'argument of perihelion',
                          'Longitude of ascending node', 'Absolute magnitude', 'Slope parameter',
                          'Orbital period in years', 'Epoch']
        if not any(keyword in data for keyword in comet_keywords):
            return False

        extracted_data = {}

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if 'C/' in line or 'P/' in line or 'SWAN' in line or 'COMET' in line.upper():
                name_match = re.search(
                    r'(C/\d+ [A-Z]+\w*|P/\d+ [A-Z]+\w*|[A-Z]+\d+ [A-Z]+\w*|Comet [A-Z0-9/ ]+)', line, re.IGNORECASE)
                if name_match:
                    extracted_data['name'] = name_match.group(1)

            if 'perihelion distance' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['q'] = match.group(1)
            elif 'semi-major axes' in line.lower() or 'semi-major axis' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['a'] = match.group(1)
            elif 'eccentricity' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['e'] = match.group(1)
            elif 'argument of perihelion' in line and '°' not in line:
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['ω'] = match.group(1)
            elif 'longitude of ascending node' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['Ω'] = match.group(1)
            elif 'inclination' in line.lower() and '°' not in line:
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['i'] = match.group(1)
            elif 'orbital period' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['P'] = match.group(1)
            elif 'time of perihelion passage' in line.lower():
                date_match = re.search(r'(\d{4}-\d{2}-\d{2}[\s\d:]+)', line)
                if date_match:
                    extracted_data['T'] = date_match.group(1).strip()
            elif 'epoch' in line.lower() and not line.lower().startswith('epoch jd'):
                epoch_match = re.search(
                    r'(\d{4}[-\s]\w+[-\s]\d+|\d{4}-\d{2}-\d{2})', line)
                if epoch_match:
                    extracted_data['epoch'] = epoch_match.group(1)
            elif 'absolute magnitude' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['H'] = match.group(1)
            elif 'slope parameter' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['n'] = match.group(1)

        if extracted_data:
            self.comet_type_var.set("彗星")
            self.clear_entries()

            if 'name' in extracted_data:
                self.name_entry.delete(0, tk.END)
                self.name_entry.insert(0, extracted_data['name'])

            orbit_mapping = {
                'q': 'perihelion_distance',
                'a': 'semi_major_axis',
                'e': 'eccentricity',
                'ω': 'arg_perihelion',
                'Ω': 'long_asc_node',
                'i': 'inclination',
                'P': 'orbital_period',
                'T': 'perihelion_time',
                'epoch': 'epoch_time'
            }

            for key, entry_key in orbit_mapping.items():
                if key in extracted_data:
                    self.orbit_entries[entry_key].delete(0, tk.END)
                    self.orbit_entries[entry_key].insert(
                        0, extracted_data[key])

            if 'H' in extracted_data and 'n' in extracted_data:
                self.cobs_calc_H_entry.delete(0, tk.END)
                self.cobs_calc_H_entry.insert(0, extracted_data['H'])
                self.cobs_calc_n_entry.delete(0, tk.END)
                self.cobs_calc_n_entry.insert(0, extracted_data['n'])
                self.use_cobs_calc_var.set(True)

            return True
        return False

    def _detect_and_import_asteroid_data(self, data):
        """检测并导入小行星轨道数据"""
        lines = data.split('\n')

        asteroid_keywords = ['perihelion distance', 'eccentricity', 'argument of perihelion',
                             'ascending node', 'inclination', 'absolute magnitude',
                             'semimajor axis', 'period', 'epoch']

        has_asteroid_features = any(keyword in data.lower()
                                    for keyword in asteroid_keywords)
        has_comet_features = any(keyword in data.lower() for keyword in [
                                 'perihelion passage', 'slope parameter', 'c/', 'p/'])

        if not has_asteroid_features or has_comet_features:
            return False

        extracted_data = {}
        asteroid_detected = False

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if '(' in line and ')' in line and any(char.isdigit() for char in line):
                name_match = re.search(r'\((\d+)\)', line)
                if name_match:
                    extracted_data['name'] = f"小行星 {name_match.group(1)}"
                    asteroid_detected = True

            if 'perihelion distance' in line.lower() and 'au' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['q'] = match.group(1)
                    asteroid_detected = True
            elif 'eccentricity' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['e'] = match.group(1)
                    asteroid_detected = True
            elif 'argument of perihelion' in line.lower() and '°' in line:
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['ω'] = match.group(1)
                    asteroid_detected = True
            elif 'ascending node' in line.lower() and '°' in line:
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['Ω'] = match.group(1)
                    asteroid_detected = True
            elif 'inclination' in line.lower() and '°' in line:
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['i'] = match.group(1)
                    asteroid_detected = True
            elif 'perihelion date' in line.lower() or 'perihelion date' in line.lower():
                date_match = re.search(r'(\d{4}-\d{2}-\d{2}\.\d+)', line)
                if date_match:
                    extracted_data['T'] = date_match.group(1)
                    asteroid_detected = True
            elif 'semimajor axis' in line.lower() and 'au' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['a'] = match.group(1)
                    asteroid_detected = True
            elif 'period' in line.lower() and 'years' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['P'] = match.group(1)
                    asteroid_detected = True
            elif 'epoch' in line.lower() and not line.lower().startswith('epoch jd'):
                epoch_match = re.search(r'(\d{4}-\d{2}-\d{2}\.\d+)', line)
                if epoch_match:
                    extracted_data['epoch'] = epoch_match.group(1)
                    asteroid_detected = True
            elif 'absolute magnitude' in line.lower():
                match = re.search(r'(\d+\.\d+)', line)
                if match:
                    extracted_data['H'] = match.group(1)
                    asteroid_detected = True

        if extracted_data and asteroid_detected:
            self.comet_type_var.set("小行星")
            self.clear_entries()

            if 'name' in extracted_data:
                self.name_entry.delete(0, tk.END)
                self.name_entry.insert(0, extracted_data['name'])
            else:
                self.name_entry.delete(0, tk.END)
                self.name_entry.insert(0, "小行星")

            orbit_mapping = {
                'q': 'perihelion_distance',
                'a': 'semi_major_axis',
                'e': 'eccentricity',
                'ω': 'arg_perihelion',
                'Ω': 'long_asc_node',
                'i': 'inclination',
                'P': 'orbital_period',
                'T': 'perihelion_time',
                'epoch': 'epoch_time'
            }

            for key, entry_key in orbit_mapping.items():
                if key in extracted_data:
                    self.orbit_entries[entry_key].delete(0, tk.END)
                    self.orbit_entries[entry_key].insert(
                        0, extracted_data[key])

            if 'H' in extracted_data:
                self.mpc_H_entry.delete(0, tk.END)
                self.mpc_H_entry.insert(0, extracted_data['H'])
                self.mpc_n_entry.delete(0, tk.END)
                self.mpc_n_entry.insert(0, "2.0")

                self.use_mpc_var.set(True)
                self.use_cobs_calc_var.set(False)
                self.use_cobs_obs_var.set(False)
                self.use_astro_var.set(False)

            return True
        return False

    def clear_import_data(self):
        """清空导入数据文本框"""
        self.data_text.delete("1.0", tk.END)

    def import_sample_data(self):
        """导入示例数据"""
        sample_data = """Comet C/2025 R2 (SWAN)
q (perihelion distance)    0.504093
a (semi-major axes)        84.678817
e (eccentricity)           0.994047
ω (argument of perihelion) 308.3520
Ω (Longitude of ascending node) 335.3310
i (inclination)            4.4728
P (Orbital period in years) 779.2551
T (time of perihelion passage) 2025-09-12 20:12:11
Epoch 2025 Nov 21
Absolute magnitude         8.9
Slope parameter            5.7"""

        self.data_text.delete("1.0", tk.END)
        self.data_text.insert("1.0", sample_data)

    # -------------------- 事件处理 --------------------
    def toggle_pred_mode(self, *_):
        if self.pred_mode.get() == "days":
            self.days_frame.pack(fill=tk.X, pady=5)
            self.dates_frame.pack_forget()
        else:
            self.days_frame.pack_forget()
            self.dates_frame.pack(fill=tk.X, pady=5)
            if not self.start_date_entry.get():
                self.start_date_entry.insert(
                    0, datetime.now().strftime('%Y-%m-%d'))
            if not self.end_date_entry.get():
                self.end_date_entry.insert(
                    0, (datetime.now() + timedelta(days=1825)).strftime('%Y-%m-%d'))

    def on_comet_type_selected(self, *_):
        self.clear_entries()
        if self.comet_type_var.get() == "彗星":
            self._build_comet_input()
        else:
            self._build_asteroid_input()

    def _build_comet_input(self):
        for key in list(self.orbit_entries.keys())[:-1]:
            self.orbit_entries[key].delete(0, tk.END)
            self.orbit_entries[key].insert(0, "")

        self.orbit_entries["perihelion_time"].delete(0, tk.END)
        self.orbit_entries["perihelion_time"].insert(0, "2025-11-08 12:54:17")

    def _build_asteroid_input(self):
        for key in list(self.orbit_entries.keys()):
            self.orbit_entries[key].delete(0, tk.END)
            self.orbit_entries[key].insert(0, "")

        self.orbit_entries["perihelion_time"].delete(0, tk.END)
        self.orbit_entries["perihelion_time"].insert(0, "2024-08-04.18975")

    # -------------------- 数据管理 --------------------
    def fill_sample_data(self):
        self.comet_type_var.set("彗星")
        self.fill_comet_data("323P-B/SOHO")

    def fill_comet_data(self, name):
        data = COMET_SAMPLE_DATA.get(name)
        if not data:
            return
        self.clear_entries()
        self.name_entry.insert(0, data["name"])
        for k, v in data["orbit"].items():
            if v:
                self.orbit_entries[k].delete(0, tk.END)
                self.orbit_entries[k].insert(0, str(v))
        self.use_cobs_calc_var.set(data["model_params"]["use_cobs_calc"])
        self.use_cobs_obs_var.set(data["model_params"]["use_cobs_obs"])
        self.use_astro_var.set(data["model_params"]["use_astro"])
        self.use_mpc_var.set(data["model_params"]["use_mpc"])
        self.cobs_calc_H_entry.insert(
            0, str(data["model_params"]["cobs_calc_H"]))
        self.cobs_calc_n_entry.insert(
            0, str(data["model_params"]["cobs_calc_n"]))
        self.cobs_obs_H_entry.insert(
            0, str(data["model_params"]["cobs_obs_H"]))
        self.cobs_obs_n_entry.insert(
            0, str(data["model_params"]["cobs_obs_n"]))
        self.astro_H_entry.insert(0, str(data["model_params"]["astro_H"]))
        self.astro_n_entry.insert(0, str(data["model_params"]["astro_n"]))
        self.mpc_H_entry.insert(0, str(data["model_params"]["mpc_H"]))
        self.mpc_n_entry.insert(0, str(data["model_params"]["mpc_n"]))

    def clear_entries(self):
        for e in self.orbit_entries.values():
            e.delete(0, tk.END)
        for h, n in [(self.cobs_calc_H_entry, self.cobs_calc_n_entry),
                     (self.cobs_obs_H_entry, self.cobs_obs_n_entry),
                     (self.astro_H_entry, self.astro_n_entry),
                     (self.mpc_H_entry, self.mpc_n_entry)]:
            h.delete(0, tk.END)
            n.delete(0, tk.END)
        for lab in self.info_labels.values():
            lab.config(text="待计算")

    # -------------------- 输入验证 --------------------
    def validate_inputs(self):
        try:
            name = self.name_entry.get().strip()
            if not name:
                raise ValueError("请输入天体名称")
            q = float(self.orbit_entries["perihelion_distance"].get())
            e = float(self.orbit_entries["eccentricity"].get())
            peri_str = self.orbit_entries["perihelion_time"].get()
            a_str = self.orbit_entries["semi_major_axis"].get()
            a = float(a_str) if a_str else self.calc_semi_major_axis(q, e)

            try:
                if self.comet_type_var.get() == "彗星":
                    peri_date = datetime.strptime(
                        peri_str, "%Y-%m-%d %H:%M:%S")
                else:
                    if '.' in peri_str:
                        peri_date = datetime.strptime(
                            peri_str.split('.')[0], "%Y-%m-%d")
                    else:
                        peri_date = datetime.strptime(peri_str, "%Y-%m-%d")
            except ValueError:
                raise ValueError(
                    "近日点时间格式错误，彗星请使用 'YYYY-MM-DD HH:MM:SS'，小行星请使用 'YYYY-MM-DD'")

            calc_H = self.cobs_calc_H_entry.get()
            calc_n = self.cobs_calc_n_entry.get()
            obs_H = self.cobs_obs_H_entry.get()
            obs_n = self.cobs_obs_n_entry.get()
            astro_H = self.astro_H_entry.get()
            astro_n = self.astro_n_entry.get()
            mpc_H = self.mpc_H_entry.get()
            mpc_n = self.mpc_n_entry.get()

            if not 0.001 <= q <= 1000:
                raise ValueError("近日点距离应在0.001-1000 AU范围内")
            if e < 0:
                raise ValueError("偏心率不能为负")
            if self.use_cobs_calc_var.get() and calc_H and calc_n:
                if not (-50 <= float(calc_H) <= 5000 and -50 <= float(calc_n) <= 5000):
                    raise ValueError("COBS-Calc模型参数应在-50到5000范围内")
            if self.use_cobs_obs_var.get() and obs_H and obs_n:
                if not (-50 <= float(obs_H) <= 5000 and -50 <= float(obs_n) <= 5000):
                    raise ValueError("COBS-Observe模型参数应在-50到5000范围内")
            if self.use_astro_var.get() and astro_H and astro_n:
                if not (-50 <= float(astro_H) <= 5000 and -50 <= float(astro_n) <= 5000):
                    raise ValueError("ASTRO模型参数应在-50到5000范围内")
            if self.use_mpc_var.get() and mpc_H and mpc_n:
                if not (-50 <= float(mpc_H) <= 5000 and -50 <= float(mpc_n) <= 5000):
                    raise ValueError("MPC模型参数应在-50到5000范围内")
            return (name, q, a, e, peri_date, calc_H, calc_n, obs_H, obs_n, astro_H, astro_n, mpc_H, mpc_n)
        except Exception as e:
            messagebox.showerror("输入错误", str(e))
            return None

    def calc_semi_major_axis(self, q, e):
        if e < 1:
            return q / (1 - e)
        elif e == 1:
            return 1e10
        else:
            return q / (1 - e)

    # -------------------- 预测计算 --------------------
    def run_prediction(self):
        inputs = self.validate_inputs()
        if inputs is None:
            return
        threading.Thread(target=self._do_prediction,
                         args=(inputs,), daemon=True).start()

    def _do_prediction(self, inputs):
        self._set_progress(0, "开始计算……")
        try:
            (name, q, a, e, peri_date, calc_H, calc_n, obs_H,
             obs_n, astro_H, astro_n, mpc_H, mpc_n) = inputs

            if self.pred_mode.get() == "days":
                future_days = int(self.pred_days_entry.get())
                if not 1 <= future_days <= 1_000_000:
                    raise ValueError("预测天数应在1-1,000,000范围内")
                start_date = peri_date - timedelta(days=future_days // 2)
                all_dates = [start_date +
                             timedelta(days=i) for i in range(future_days)]
            else:
                start_date = datetime.strptime(
                    self.start_date_entry.get(), '%Y-%m-%d')
                end_date = datetime.strptime(
                    self.end_date_entry.get(), '%Y-%m-%d')
                delta_days = (end_date - start_date).days
                if delta_days > 1_000_000:
                    raise ValueError("预测范围不能超过1,000,000天")
                all_dates = [start_date +
                             timedelta(days=i) for i in range(delta_days + 1)]

            total = len(all_dates)
            self._set_progress(5, "计算日彗距/地彗距……")

            # 计算日彗距
            r_pred = []
            for i, date in enumerate(all_dates):
                days_from_peri = (date - peri_date).days
                if e < 1:
                    M = 2 * np.pi * days_from_peri / (a ** 1.5 * 365.25)
                    E = self.solve_kepler(M, e)
                    r = a * (1 - e * np.cos(E))
                elif e == 1:
                    r = self.calc_parabolic(days_from_peri, q)
                else:
                    r = self.calc_hyperbolic(days_from_peri, q, e, a)
                r_pred.append(max(r, 0.001))
                if i % 100 == 0:
                    self._set_progress(5 + 25 * i // total, "计算日彗距/地彗距……")
            r_pred = np.array(r_pred)

            # 计算地彗距
            delta_pred = []
            for i, date in enumerate(all_dates):
                delta = self.calc_earth_distance(date, r_pred[i], e, peri_date)
                delta_pred.append(delta)
                if i % 100 == 0:
                    self._set_progress(30 + 20 * i // total, "计算地彗距……")
            delta_pred = np.array(delta_pred)

            # 计算星等
            self._set_progress(55, "计算星等……")
            self.cobs_calc_m_pred = None
            self.cobs_obs_m_pred = None
            self.astro_m_pred = None
            self.mpc_m_pred = None

            if self.use_cobs_calc_var.get() and calc_H and calc_n:
                self.cobs_calc_m_pred = (float(
                    calc_H) + 5 * np.log10(delta_pred) + 2.5 * float(calc_n) * np.log10(r_pred))
                self.cobs_calc_m_pred = np.clip(
                    self.cobs_calc_m_pred, -100, 5000)
            if self.use_cobs_obs_var.get() and obs_H and obs_n:
                self.cobs_obs_m_pred = float(
                    obs_H) + 5 * np.log10(delta_pred) + 2.5 * float(obs_n) * np.log10(r_pred)
                self.cobs_obs_m_pred = np.clip(
                    self.cobs_obs_m_pred, -100, 5000)
            if self.use_astro_var.get() and astro_H and astro_n:
                self.astro_m_pred = float(
                    astro_H) + 5 * np.log10(delta_pred) + 2.5 * float(astro_n) * np.log10(r_pred)
                self.astro_m_pred = np.clip(self.astro_m_pred, -100, 5000)
            if self.use_mpc_var.get() and mpc_H and mpc_n:
                self.mpc_m_pred = float(
                    mpc_H) + 5 * np.log10(delta_pred) + 2.5 * float(mpc_n) * np.log10(r_pred)
                self.mpc_m_pred = np.clip(self.mpc_m_pred, -100, 5000)

            self.current_dates = all_dates
            self.r_pred = r_pred
            self.delta_pred = delta_pred

            self._set_progress(90, "更新界面……")
            self.update_results(name, e, all_dates, r_pred, total)
            self.update_brightness_window()
            self.update_distance_window()
            self._set_progress(100, "完成")
            messagebox.showinfo("预测完成", "天体亮度预测计算完成！ (v8.1)")

        except Exception as e:
            self._set_progress(0, "错误")
            messagebox.showerror("预测错误", str(e))

    def _set_progress(self, val, text):
        self.progress_var.set(val)
        self.progress_text.set(text)
        self.root.update_idletasks()

    # -------------------- 轨道算法 --------------------
    def solve_kepler(self, M, e, tol=1e-8, max_iter=100):
        if abs(M) < 1e-6:
            return M
        E = M if e < 0.8 else (np.pi if M > 0 else -np.pi)
        for _ in range(max_iter):
            delta_E = (E - e * np.sin(E) - M) / (1 - e * np.cos(E))
            E -= delta_E
            if abs(E) > 10 * np.pi:
                E = np.sign(E) * 10 * np.pi
            if abs(delta_E) < tol:
                break
        return E

    def calc_parabolic(self, days_from_peri, q):
        W = 3 * days_from_peri * np.sqrt(0.5 / (q ** 3)) / 2
        s = W
        for _ in range(10):
            s_new = W - s ** 3 / 3
            if abs(s_new - s) < 1e-8:
                break
            s = s_new
        return max(q * (1 + s ** 2), 0.001)

    def calc_hyperbolic(self, days_from_peri, q, e, a):
        N = 2 * np.pi * days_from_peri / ((-a) ** 1.5 * 365.25)
        H = self.solve_hyperbolic_kepler(N, e)
        return max(a * (1 - e * np.cosh(H)), 0.001)

    def solve_hyperbolic_kepler(self, N, e, tol=1e-8, max_iter=100):
        if abs(N) < 0.1:
            H = N
        else:
            H = np.sign(N) * np.log(2 * abs(N) / e + 1.8)
        for _ in range(max_iter):
            sinhH, coshH = np.sinh(H), np.cosh(H)
            f = e * sinhH - H - N
            df = e * coshH - 1
            if abs(df) < 1e-12:
                H += 0.1
                continue
            delta = f / df
            H -= delta
            if abs(H) > 100:
                H = np.sign(H) * 100
            if abs(delta) < tol:
                break
        return H

    def calc_earth_distance(self, date, r, e, peri_date):
        days_from_peri = (date - peri_date).days
        earth_angle = 2 * np.pi * (days_from_peri % 365.25) / 365.25
        earth_x, earth_y = np.cos(earth_angle), np.sin(earth_angle)
        omega = np.radians(float(self.orbit_entries["arg_perihelion"].get()))
        i = np.radians(float(self.orbit_entries["inclination"].get()))
        comet_x = r * np.cos(omega)
        comet_y = r * np.sin(omega) * np.cos(i)
        delta = np.sqrt((comet_x - earth_x) ** 2 + (comet_y - earth_y) ** 2)
        return max(delta, 0.001)

    # -------------------- 结果更新 --------------------
    def update_results(self, name, e, dates, r_pred, total):
        orbit_type = f"椭圆轨道 (e={e:.6f})" if e < 1 else "抛物线轨道 (e=1.000000)" if e == 1 else f"双曲线轨道 (e={e:.6f})"
        self.info_labels["轨道类型"].config(text=orbit_type)
        self.info_labels["预测总天数"].config(text=f"{total:,} 天")

        peri_idx = np.argmin(r_pred)
        peri_date_str = dates[peri_idx].strftime('%Y-%m-%d')
        self.info_labels["近日点日期"].config(text=peri_date_str)
        self.info_labels["近地点距离"].config(text=f"{r_pred[peri_idx]:.7f} AU")

        today = datetime.now().date()
        today_idx = next((i for i, date in enumerate(
            dates) if date.date() == today), None)
        if today_idx is not None:
            today_data = {
                "calc": self.cobs_calc_m_pred[today_idx] if self.cobs_calc_m_pred is not None else "N/A",
                "obs": self.cobs_obs_m_pred[today_idx] if self.cobs_obs_m_pred is not None else "N/A",
                "astro": self.astro_m_pred[today_idx] if self.astro_m_pred is not None else "N/A",
                "mpc": self.mpc_m_pred[today_idx] if self.mpc_m_pred is not None else "N/A",
                "r": r_pred[today_idx],
                "delta": self.delta_pred[today_idx]
            }
            self.info_labels["今日日彗距"].config(text=f"{today_data['r']:.7f}")
            self.info_labels["今日地彗距"].config(
                text=f"{today_data['delta']:.7f} AU")
            for k, lbl in zip(["calc", "obs", "astro", "mpc"], ["今日COBS-Calc星等", "今日COBS-Observe星等", "今日ASTRO星等", "今日MPC星等"]):
                if today_data[k] != "N/A":
                    self.info_labels[lbl].config(text=f"{today_data[k]:.2f}")

        for days in [30, 60]:
            future_date = datetime.now() + timedelta(days=days)
            future_idx = next((i for i, date in enumerate(
                dates) if date.date() == future_date.date()), None)
            if future_idx is not None:
                future_data = {
                    "calc": self.cobs_calc_m_pred[future_idx] if self.cobs_calc_m_pred is not None else "N/A",
                    "obs": self.cobs_obs_m_pred[future_idx] if self.cobs_obs_m_pred is not None else "N/A",
                    "astro": self.astro_m_pred[future_idx] if self.astro_m_pred is not None else "N/A",
                    "mpc": self.mpc_m_pred[future_idx] if self.mpc_m_pred is not None else "N/A",
                    "r": r_pred[future_idx],
                    "delta": self.delta_pred[future_idx]
                }
                self.info_labels[f"{days}天后日彗距"].config(
                    text=f"{future_data['r']:.7f}")
                self.info_labels[f"{days}天后地彗距"].config(
                    text=f"{future_data['delta']:.7f} AU")
                for k, lbl in zip(["calc", "obs", "astro", "mpc"], [f"{days}天后COBS-Calc星等", f"{days}天后COBS-Observe星等", f"{days}天后ASTRO星等", f"{days}天后MPC星等"]):
                    if future_data[k] != "N/A":
                        self.info_labels[lbl].config(
                            text=f"{future_data[k]:.2f}")

        if today_idx is not None:
            coma_length_km, coma_length_deg, tail_length_km, tail_length_deg = self.calc_coma_tail(
                r_pred[today_idx], self.delta_pred[today_idx])
            self.info_labels["彗发长度 (km)"].config(text=f"{coma_length_km:.0f}")
            self.info_labels["彗发长度 (°)"].config(text=f"{coma_length_deg:.4f}")
            self.info_labels["彗尾长度 (km)"].config(text=f"{tail_length_km:.0f}")
            self.info_labels["彗尾长度 (°)"].config(text=f"{tail_length_deg:.4f}")

    def calc_coma_tail(self, r, delta):
        coma_length_km = 100000 / np.sqrt(r)
        tail_length_km = 1000000 / np.sqrt(r)
        delta_km = delta * 149597870.7
        coma_length_deg = np.degrees(coma_length_km / delta_km)
        tail_length_deg = np.degrees(tail_length_km / delta_km)
        return coma_length_km, coma_length_deg, tail_length_km, tail_length_deg

    # -------------------- 图表窗口 --------------------
    def open_brightness_window(self):
        if self.brightness_window is None or not self.brightness_window.winfo_exists():
            self.brightness_window = tk.Toplevel(self.root)
            self.brightness_window.title("亮度预测图表 v8.1")
            self.brightness_window.geometry("1000x700")
            self.setup_brightness_window()
            if self.current_dates is not None:
                self.update_brightness_window()
        else:
            self.brightness_window.lift()

    def open_distance_window(self):
        if self.distance_window is None or not self.distance_window.winfo_exists():
            self.distance_window = tk.Toplevel(self.root)
            self.distance_window.title("距离变化图表 v8.1")
            self.distance_window.geometry("1000x700")
            self.setup_distance_window()
            if self.current_dates is not None:
                self.update_distance_window()
        else:
            self.distance_window.lift()

    def close_all_windows(self):
        for w in [self.brightness_window, self.distance_window]:
            if w and w.winfo_exists():
                w.destroy()

    def setup_brightness_window(self):
        if not self.brightness_window.winfo_exists():
            return
        self.brightness_fig, self.brightness_ax = plt.subplots(figsize=(12, 8))
        self.brightness_canvas = FigureCanvasTkAgg(
            self.brightness_fig, master=self.brightness_window)
        self.brightness_canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        NavigationToolbar2Tk(self.brightness_canvas,
                             self.brightness_window).update()

    def setup_distance_window(self):
        if not self.distance_window.winfo_exists():
            return
        self.distance_fig, self.distance_ax = plt.subplots(figsize=(12, 8))
        self.distance_canvas = FigureCanvasTkAgg(
            self.distance_fig, master=self.distance_window)
        self.distance_canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        NavigationToolbar2Tk(self.distance_canvas,
                             self.distance_window).update()

    def update_brightness_window(self):
        if (self.brightness_window is None or not self.brightness_window.winfo_exists() or self.current_dates is None):
            return
        self.brightness_ax.clear()
        e = float(self.orbit_entries["eccentricity"].get())
        orbit_type = "椭圆轨道" if e < 1 else ("抛物线轨道" if e == 1 else "双曲线轨道")
        colors = ['#d62728', '#2ca02c', '#000080', '#ff7f0e']
        labels = ['COBS-Calc', 'COBS-Observe', 'ASTRO', 'MPC']
        predictions = [self.cobs_calc_m_pred,
                       self.cobs_obs_m_pred, self.astro_m_pred, self.mpc_m_pred]
        use_models = [self.use_cobs_calc_var.get(), self.use_cobs_obs_var.get(
        ), self.use_astro_var.get(), self.use_mpc_var.get()]

        # 绘制预测曲线
        for i, (pred, color, label, use) in enumerate(zip(predictions, colors, labels, use_models)):
            if pred is not None and use:
                self.brightness_ax.plot(
                    self.current_dates, pred, color=color, label=label, linewidth=2)
                min_mag = min(pred)
                min_idx = np.argmin(pred)
                self.brightness_ax.plot(
                    self.current_dates[min_idx], min_mag, 'o', color=color, markersize=8, label=f'{label}最亮: {min_mag:.2f}等')

        # 绘制观测数据
        current_type = self.comet_type_var.get()
        if current_type == "彗星":
            obs_data = self.observation_data['comet']
            obs_color = 'red'
            obs_label = '彗星观测'
        else:
            obs_data = self.observation_data['asteroid']
            obs_color = 'blue'
            obs_label = '小行星观测'

        if obs_data:
            pred_start_date = min(self.current_dates)
            pred_end_date = max(self.current_dates)
            filtered_obs = []
            for obs in obs_data:
                if pred_start_date <= obs['date'] <= pred_end_date:
                    filtered_obs.append(obs)

            if filtered_obs:
                obs_dates = [obs['date'] for obs in filtered_obs]
                obs_mags = [obs['magnitude'] for obs in filtered_obs]
                self.brightness_ax.scatter(
                    obs_dates, obs_mags, color=obs_color, label=obs_label, s=20, alpha=0.7, zorder=5)

        # 标记今日和30天后
        today = datetime.now().date()
        today_idx = next((i for i, date in enumerate(
            self.current_dates) if date.date() == today), None)
        if today_idx is not None:
            for i, (pred, color, label, use) in enumerate(zip(predictions, colors, labels, use_models)):
                if pred is not None and use:
                    today_mag = pred[today_idx]
                    self.brightness_ax.plot(
                        self.current_dates[today_idx], today_mag, 's', color=color, markersize=6, label=f'今日{label}: {today_mag:.2f}等')

        thirty_days_later = datetime.now() + timedelta(days=30)
        thirty_idx = next((i for i, date in enumerate(
            self.current_dates) if date.date() == thirty_days_later.date()), None)
        if thirty_idx is not None:
            for i, (pred, color, label, use) in enumerate(zip(predictions, colors, labels, use_models)):
                if pred is not None and use:
                    thirty_mag = pred[thirty_idx]
                    self.brightness_ax.plot(
                        self.current_dates[thirty_idx], thirty_mag, 'D', color=color, markersize=6, label=f'30天后{label}: {thirty_mag:.2f}等')

        self.brightness_ax.set_ylabel("视星等")
        self.brightness_ax.set_title(
            f"{self.name_entry.get()} 亮度预测 - {orbit_type} (v8.1)", fontsize=14)
        self.brightness_ax.grid(True, alpha=0.3)
        self.brightness_ax.legend()
        self.brightness_ax.invert_yaxis()
        self.brightness_fig.autofmt_xdate()
        self.brightness_fig.tight_layout()
        self.brightness_canvas.draw()

    def update_distance_window(self):
        if (self.distance_window is None or not self.distance_window.winfo_exists() or self.current_dates is None):
            return
        self.distance_ax.clear()
        e = float(self.orbit_entries["eccentricity"].get())
        orbit_type = "椭圆轨道" if e < 1 else ("抛物线轨道" if e == 1 else "双曲线轨道")
        self.distance_ax.plot(self.current_dates, self.r_pred,
                              'orange', label='日彗距', linewidth=2)
        self.distance_ax.plot(
            self.current_dates, self.delta_pred, 'blue', label='地彗距', linewidth=2)

        today_idx = next((i for i, date in enumerate(
            self.current_dates) if date.date() == datetime.now().date()), None)
        if today_idx is not None:
            today_r, today_delta = self.r_pred[today_idx], self.delta_pred[today_idx]
            self.distance_ax.plot(self.current_dates[today_idx], today_r, 's',
                                  color='purple', markersize=6, label=f'今日日彗距: {today_r:.7f} AU')
            self.distance_ax.plot(self.current_dates[today_idx], today_delta, 's',
                                  color='green', markersize=6, label=f'今日地彗距: {today_delta:.7f} AU')

        peri_date_str = self.orbit_entries["perihelion_time"].get()
        if peri_date_str:
            try:
                if self.comet_type_var.get() == "彗星":
                    peri_date = datetime.strptime(
                        peri_date_str, "%Y-%m-%d %H:%M:%S")
                else:
                    peri_date = datetime.strptime(
                        peri_date_str.split('.')[0], "%Y-%m-%d")
                peri_idx = next((i for i, date in enumerate(
                    self.current_dates) if date.date() == peri_date.date()), None)
                if peri_idx is not None:
                    peri_r, peri_delta = self.r_pred[peri_idx], self.delta_pred[peri_idx]
                    self.distance_ax.plot(
                        self.current_dates[peri_idx], peri_r, 'D', color='orange', markersize=6, label=f'近日点: {peri_r:.7f} AU')
                    self.distance_ax.plot(
                        self.current_dates[peri_idx], peri_delta, 'D', color='cyan', markersize=6, label=f'近日点地彗距: {peri_delta:.7f} AU')
            except ValueError:
                pass

        self.distance_ax.set_ylabel("距离 (AU)")
        self.distance_ax.set_xlabel("日期")
        self.distance_ax.set_title(
            f"{self.name_entry.get()} 距离变化 - {orbit_type} (v8.1)", fontsize=14)
        self.distance_ax.grid(True, alpha=0.3)
        self.distance_ax.legend()
        self.distance_fig.autofmt_xdate()
        self.distance_fig.tight_layout()
        self.distance_canvas.draw()


# -------------------- 样本数据 --------------------
COMET_SAMPLE_DATA = {
    "323P-B/SOHO": {
        "name": "323P-B/SOHO",
        "orbit": {
            "perihelion_distance": "0.040063",
            "semi_major_axis": "2.890131",
            "eccentricity": "0.986138",
            "arg_perihelion": "353.9691",
            "long_asc_node": "323.4650",
            "inclination": "5.4627",
            "orbital_period": "4.9135",
            "perihelion_time": "2025-12-16 05:37:06",
            "epoch_time": "2025-10-06"
        },
        "model_params": {
            "use_cobs_calc": True,
            "use_cobs_obs": True,
            "use_astro": True,
            "use_mpc": True,
            "cobs_calc_H": "20.0",
            "cobs_calc_n": "10.0",
            "cobs_obs_H": "19.5",
            "cobs_obs_n": "9.5",
            "astro_H": "18.5",
            "astro_n": "9.0",
            "mpc_H": "21.0",
            "mpc_n": "5.0"
        }
    }
}

# -------------------- 启动 --------------------
if __name__ == "__main__":
    root = tk.Tk()
    app = CometBrightnessApp(root)
    root.mainloop()
